﻿using System;

namespace SkeinLibManaged
{   //Pseudorandom Number Generator (PRNG)

    public partial class Skein
    {   //STATE VARIABLES
        internal UInt64[] PRNGState256 = null;
        internal UInt64[] PRNGState512 = null;
        internal UInt64[] PRNGState1K = null;

        /// <summary>
        /// Is the PRNG initialized?
        /// </summary>
        public bool PRNGInitialized(StateSize which)
        {
            switch (which)
            {
                case StateSize.Narrow: return (PRNGState256 != null);
                case StateSize.Medium: return (PRNGState512 != null);
                case StateSize.Wide: return (PRNGState1K != null);
                default: return false;
            }
        }

        /// <summary>
        /// Initialize a PRNG state with a buffer of data."
        /// </summary>
        /// <param name="stateBitWidth">Enum of which PRNG to update.</param>
        /// <param name="seed">The buffer holding the data into which to PRNG will be initialized. </param>
        /// <returns>DEBUG ONLY: Returns the updated state buffer.</returns>
        public UInt64[] InitializePRNG(StateSize stateBitWidth, byte[] seed)
        {
            if (seed == null || seed.Length == 0)
                throw new ArgumentOutOfRangeException("InitializePRNG:  Seed buffer cannot be null or zero length!");
            UInt64[] tweak = { 0, 0 };
            UInt16 BYTEWIDTH = (UInt16)(UInt16.Parse(System.Enum.Format(typeof(StateSize), stateBitWidth, "D")) / 8);
            UInt16 WORDWIDTH = (UInt16)(BYTEWIDTH / sizeof(UInt64));
            UInt64[] state = new UInt64[WORDWIDTH];
            UInt64[] transport = new UInt64[WORDWIDTH];
            ClearWords(state);

            //BEGIN CONFIG
            ClearWords(transport);
            transport[0] = SKEIN_VS;
            transport[1] = ulong.MaxValue;

            tweak[0] = 32;
            tweak[1] = SKEIN_FLAG_FIRST + SKEIN_FLAG_FINAL + SKEIN_TYPE_CONFIG;

            switch (stateBitWidth)
            {
                case StateSize.Narrow: state = ThreeFish.ThreeFish256Words(transport, state, tweak, true); break;
                case StateSize.Medium: state = ThreeFish.ThreeFish512Words(transport, state, tweak, true); break;
                case StateSize.Wide: state = ThreeFish.ThreeFish1024Words(transport, state, tweak, true); break;
            }
            //END CONFIG

            //BEGIN SEED
            tweak[0] = 0;
            tweak[1] = SKEIN_FLAG_FIRST + SKEIN_TYPE_MSG;
            for (long i = 0; (i * BYTEWIDTH) < seed.Length; i++)
            {
                //process a block
                tweak[0] = (UInt64)(i * BYTEWIDTH);
                byte[] buf = new byte[BYTEWIDTH];
                for (int j = 0; j < BYTEWIDTH; j++)
                {
                    //process a byte
                    if ((long)(j + (i * BYTEWIDTH)) < seed.Length)
                    {
                        buf[j] = seed[(i * BYTEWIDTH) + j];
                        tweak[0] += 1;  //count each byte
                    }
                    else	// or zero it out
                        buf[j] = 0;
                }
                if ((i + 1) * BYTEWIDTH >= seed.Length)
                    tweak[1] += SKEIN_FLAG_FINAL;
                //stuff transport in prep for transformation
                System.Buffer.BlockCopy(buf, 0, transport, 0, buf.Length);
                switch (stateBitWidth)
                {
                    case StateSize.Narrow: state = ThreeFish.ThreeFish256Words(transport, state, tweak, true); break;
                    case StateSize.Medium: state = ThreeFish.ThreeFish512Words(transport, state, tweak, true); break;
                    case StateSize.Wide: state = ThreeFish.ThreeFish1024Words(transport, state, tweak, true); break;
                }
                tweak[1] &= ~SKEIN_FLAG_FIRST;
            }
            //END SEED

            //BEGIN OUTPUT
            tweak[0] = 8; // <<<<<<<<<<<<<<WHY???????? not explained well in the white paper or ref code
            tweak[1] = SKEIN_FLAG_FIRST + SKEIN_TYPE_OUTPUT + SKEIN_FLAG_FINAL;       
            ClearWords(transport);
            switch (stateBitWidth)
            {
                //transform then store to the appropriate PRNGState
                case StateSize.Narrow:
                    transport = ThreeFish.ThreeFish256Words(transport, state, tweak, true);
                    if (this.PRNGState256 == null) 
                        PRNGState256 = new UInt64[WORDWIDTH];
                    System.Buffer.BlockCopy(transport, 0, this.PRNGState256, 0, (int)BYTEWIDTH);
                    break;

                case StateSize.Medium: 
                    transport = ThreeFish.ThreeFish512Words(transport, state, tweak, true);
                    if (this.PRNGState512 == null) 
                        PRNGState512 = new UInt64[WORDWIDTH];
                    System.Buffer.BlockCopy(transport, 0, this.PRNGState512, 0, (int)BYTEWIDTH);
                    break;

                case StateSize.Wide: 
                    transport = ThreeFish.ThreeFish1024Words(transport, state, tweak, true);
                    if (this.PRNGState1K == null) 
                        PRNGState1K = new UInt64[WORDWIDTH];
                    System.Buffer.BlockCopy(transport, 0, this.PRNGState1K, 0, (int)BYTEWIDTH);
                    break;
            }
            //zero out used memory
            //TODO:  When not in debug mode, REMOVE the return transport, and clear transport
            //ClearWords(transport);
            ClearWords(state);
            ClearWords(tweak);
            //END OUTPUT

            return transport;  //<-- DEBUG ONLY!!!
        }//end InitializePRNG

        /// <summary>
        /// Update the PRNG state with an additional amount of data.  The PRNG must have already been initialized!
        /// </summary>
        /// <param name="stateBitWidth">Enum of which PRNG to update.</param>
        /// <param name="reseed">The buffer holding the data from which the PRNG will be injected. </param>
        /// <returns>DEBUG ONLY: Returns the updated state buffer.</returns>
        public UInt64[] UpdatePRNG(StateSize stateBitWidth, byte[] reseed)
        {   
            if (reseed == null || reseed.Length == 0)
                throw new ArgumentOutOfRangeException("UpdatePRNG:  Reseed buffer cannot be null or zero length!");
            UInt16 BYTEWIDTH = (UInt16)(UInt16.Parse(System.Enum.Format(typeof(StateSize), stateBitWidth, "D")) / 8);
            UInt16 WORDWIDTH = (UInt16)(BYTEWIDTH / sizeof(UInt64));
            UInt64[] state = new UInt64[WORDWIDTH];
            switch (stateBitWidth)
            {    case StateSize.Narrow:
                    if (PRNGState256 == null)
                    {
                        throw new System.NullReferenceException("UpdatePRNG:  Cannot update: " +
                            "The Narrow state has not been initialized! Call InitializePRNG first.");
                    }
                    else
                        System.Buffer.BlockCopy(PRNGState256, 0, state, 0, (int)BYTEWIDTH);
                    break;

                case StateSize.Medium:
                    if (PRNGState512 == null)
                    {
                        throw new System.NullReferenceException("UpdatePRNG:  Cannot update: " +
                            "The Medium state has not been initialized! Call InitializePRNG first.");
                    }
                    else
                        System.Buffer.BlockCopy(PRNGState512, 0, state, 0, (int)BYTEWIDTH);
                    break;

                case StateSize.Wide:
                    if (PRNGState1K == null)
                    {
                        throw new System.NullReferenceException("UpdatePRNG:  Cannot update: " +
                            "The Wide state has not been initialized! Call InitializePRNG first.");
                    }
                    else
                        System.Buffer.BlockCopy(PRNGState1K, 0, state, 0, (int)BYTEWIDTH);
                    break;

            }
            UInt64[] tweak = { 0, 0 };
            UInt64[] transport = new UInt64[WORDWIDTH];

            //BEGIN MSG
            tweak[0] = 0;
            tweak[1] = SKEIN_FLAG_FIRST + SKEIN_TYPE_MSG;
            for (long i = 0; (i * BYTEWIDTH) < reseed.Length; i++)
            {
                //process a block
                tweak[0] = (UInt64)(i * BYTEWIDTH);
                byte[] buf = new byte[BYTEWIDTH];
                for (int j = 0; j < BYTEWIDTH; j++)
                {
                    //process a byte
                    if ((long)(j + (i * BYTEWIDTH)) < reseed.Length)
                    {
                        buf[j] = reseed[(i * BYTEWIDTH) + j];
                        tweak[0] += 1;  //count each byte
                    }
                    else	// or zero it out
                        buf[j] = 0;
                }
                if ((i + 1) * BYTEWIDTH >= reseed.Length)
                    tweak[1] += SKEIN_FLAG_FINAL;
                //stuff transport in prep for transformation
                System.Buffer.BlockCopy(buf, 0, transport, 0, buf.Length);
                switch (stateBitWidth)
                {
                    case StateSize.Narrow: state = ThreeFish.ThreeFish256Words(transport, state, tweak, true); break;
                    case StateSize.Medium: state = ThreeFish.ThreeFish512Words(transport, state, tweak, true); break;
                    case StateSize.Wide: state = ThreeFish.ThreeFish1024Words(transport, state, tweak, true); break;
                }
                tweak[1] &= ~SKEIN_FLAG_FIRST;
            }
            //END MSG

            //BEGIN OUTPUT
            tweak[0] = 8; // <<<<<<<<<<<<<<WHY???????? not explained well in the white paper or ref code
            tweak[1] = SKEIN_FLAG_FIRST + SKEIN_TYPE_OUTPUT + SKEIN_FLAG_FINAL;
            switch (stateBitWidth)
            {
                //transform then store to the appropriate PRNGState
                //TODO:  When not in debug mode, move PRNGStateX variables to assigment statements
                //PRNGStateX = ThreeFish.ThreeFishXWords(transport, state, tweak, true);
                case StateSize.Narrow:
                    transport = ThreeFish.ThreeFish256Words(transport, state, tweak, true);
                    System.Buffer.BlockCopy(transport, 0, this.PRNGState256, 0, (int)BYTEWIDTH);
                    break;

                case StateSize.Medium:
                    transport = ThreeFish.ThreeFish512Words(transport, state, tweak, true);
                    System.Buffer.BlockCopy(transport, 0, this.PRNGState512, 0, (int)BYTEWIDTH);
                    break;

                case StateSize.Wide:
                    transport = ThreeFish.ThreeFish1024Words(transport, state, tweak, true);
                    System.Buffer.BlockCopy(transport, 0, this.PRNGState1K, 0, (int)BYTEWIDTH);
                    break;
            }
            //zero out used memory
            //TODO:  When not in debug mode, REMOVE the "return transport", and uncomment clearwords transport
            //ClearWords(transport);
            ClearWords(state);
            ClearWords(tweak);
            //END OUTPUT

            return transport; //<-- DEBUG ONLY!!!
        }// end UpdatePRNG

        /// <summary>
        /// Get a buffer of bytes from the PRNG.  The state is scrambled after every request, so a
        /// request of 1 byte, followed by a request of 15 bytes, will not produce the same as a
        /// single request of 16 bytes.
        /// </summary>
        /// <param name="stateBitWidth">Enum of which PRNG will be used to gather bytes.</param>
        /// <param name="byteCount">The count, in bytes, of the resulting buffer.  Max = 65536 bytes in one call.</param>
        /// <returns>A byte buffer of pseudo-random bytes of length "byteCount".</returns>
        public byte[] GetBytesFromPRNG(StateSize stateBitWidth, UInt16 byteCount)
        {
            if (byteCount == 0)
                throw new System.ArgumentOutOfRangeException("GetBytesFromPRNG:  Byte count cannot be 0! "+
                                                             "You must request at least 1 byte.");
            UInt64[] tweak = { 0, 0 };
            UInt16 BYTEWIDTH = (UInt16)(UInt16.Parse(System.Enum.Format(typeof(StateSize), stateBitWidth, "D")) / 8);
            UInt16 WORDWIDTH = (UInt16)(BYTEWIDTH / sizeof(UInt64));
            UInt64[] state = new UInt64[WORDWIDTH];
            UInt64[] transport = new UInt64[WORDWIDTH];
            switch (stateBitWidth)
            {
                case StateSize.Narrow:
                    if (PRNGState256 == null)
                    {
                        throw new System.NullReferenceException("GetBytesFromPRNG:  Cannot get bytes: " +
                            "The Narrow state has not been initialized!  Call InitializePRNG first.");
                    }
                    else
                        System.Buffer.BlockCopy(PRNGState256, 0, state, 0, (int)BYTEWIDTH);
                    break;

                case StateSize.Medium:
                    if (PRNGState512 == null)
                    {
                        throw new System.NullReferenceException("GetBytesFromPRNG:  Cannot get bytes: " +
                            "The Medium state has not been initialized!  Call InitializePRNG first.");
                    }
                    else
                        System.Buffer.BlockCopy(PRNGState512, 0, state, 0, (int)BYTEWIDTH);
                    break;

                case StateSize.Wide:
                    if (PRNGState1K == null)
                    {
                        throw new System.NullReferenceException("GetBytesFromPRNG:  Cannot get bytes: " +
                            "The Wide state has not been initialized!  Call InitializePRNG first.");
                    }
                    else
                        System.Buffer.BlockCopy(PRNGState1K, 0, state, 0, (int)BYTEWIDTH);
                    break;
            }

            //BEGIN OUTPUT
            tweak[0] = 8; // <<<<<<<<<<<<<<WHY???????? not explained well in the white paper or ref code
            tweak[1] = SKEIN_FLAG_FIRST + SKEIN_TYPE_OUTPUT + SKEIN_FLAG_FINAL;
            //consume first block updating state
            ClearWords(transport);
            switch (stateBitWidth)
            {
                case StateSize.Narrow: 
                    state = ThreeFish.ThreeFish256Words(transport, state, tweak, true);
                    System.Buffer.BlockCopy(state, 0, PRNGState256, 0, (int)BYTEWIDTH); break;
                    break;
                case StateSize.Medium: 
                    state = ThreeFish.ThreeFish512Words(transport, state, tweak, true);
                    System.Buffer.BlockCopy(state, 0, PRNGState512, 0, (int)BYTEWIDTH); break;
                    break;
                case StateSize.Wide: 
                    state = ThreeFish.ThreeFish1024Words(transport, state, tweak, true);
                    System.Buffer.BlockCopy(state, 0, PRNGState1K, 0, (int)BYTEWIDTH); break;
                    break;
            }
            //begin output of actual data
            byte[] result = new byte[byteCount];
            for (Int32 i = 0; (i * BYTEWIDTH) < byteCount; i++)
            {   //use the input array to the ThreeFish cipher for a counter, the state and tweak remain constant
                ClearWords(transport);
                transport[0] = (UInt64)(i + 1); //counter is i+1 because we already used 1 block to update the state above
                switch (stateBitWidth)
                {
                    case StateSize.Narrow: transport = ThreeFish.ThreeFish256Words(transport, state, tweak, true); break;
                    case StateSize.Medium: transport = ThreeFish.ThreeFish512Words(transport, state, tweak, true); break;
                    case StateSize.Wide: transport = ThreeFish.ThreeFish1024Words(transport, state, tweak, true); break;
                }

                //dump the transport to the output, then discard
                int position = (i * BYTEWIDTH);
                if (position <= byteCount)
                    System.Buffer.BlockCopy(transport, 0, result, position, (int)BYTEWIDTH);
                else
                    System.Buffer.BlockCopy(transport, 0, result, position, (int)(result.Length - position));
            }
            //zero out used memory
            ClearWords(transport);
            ClearWords(state);
            ClearWords(tweak);
            //END OUTPUT

            return result;
        }
    } //end class
}//end namespace